import numpy as np
import cv2

# 已知正方形的尺寸
square_size = 200 # mm
carb_xml = 'main_Para.xml'
# carb_xml = 'side_F_Para.xml'
# carb_xml = 'side_L_Para.xml'
fs = cv2.FileStorage(carb_xml, cv2.FILE_STORAGE_READ)


# 相机内参矩阵
# K = np.array([[3450, 0, 2100],
#               [0, 3450, 1080],
#               [0, 0, 1]])
K = fs.getNode('camera-matrix').mat()

#像素分辨率
pix_rev = 0.125
error = pix_rev/1.414
# 在像素坐标系下的边长
pix_lenth = 5120/3
# 正方形在像素坐标系下的四个角点坐标
# p1 = np.array([2100 - pix_lenth/2 - error, 1080 - pix_lenth/2 + error])
# p2 = np.array([2100 - pix_lenth/2 + error, 1080 + pix_lenth/2 - error])
# p3 = np.array([2100 + pix_lenth/2 - error, 1080 + pix_lenth/2 - error])
# p4 = np.array([2100 + pix_lenth/2 + error, 1080 - pix_lenth/2 + error])
p1 = np.array([K[0][2] - pix_lenth/2 - error, K[1][2] - pix_lenth/2 + error])
p2 = np.array([K[0][2] - pix_lenth/2 + error, K[1][2] + pix_lenth/2 - error])
p3 = np.array([K[0][2] + pix_lenth/2 + error, K[1][2] + pix_lenth/2 - error])
p4 = np.array([K[0][2] + pix_lenth/2 - error, K[1][2] - pix_lenth/2 - error])
# p1 = np.array([2100 - 900, 1080 - 900])
# p2 = np.array([2100 - 900, 1080 + 900])
# p3 = np.array([2100 + 900, 1080 + 900])
# p4 = np.array([2100 + 900, 1080 - 900])
# p1 = np.array([1201.319, 181.319])
# p2 = np.array([1198.681, 1978.681])
# p3 = np.array([3001.319, 1978.681])
# p4 = np.array([2998.681, 181.319])

image_points = np.array([p1, p2, p3, p4])

# 旋转矩阵变换为欧拉角
def rotation_matrix_to_euler_angles(R):
    # 计算旋转向量
    rotation_vector, _ = cv2.Rodrigues(R)
    # 构造投影矩阵
    P = np.hstack((R, np.zeros((3, 1))))
    # 分解投影矩阵以获得欧拉角
    _, _, _, _, _, _, euler_angles = cv2.decomposeProjectionMatrix(P)
    return euler_angles

# 构建正方形的三维坐标矩阵
object_points = np.zeros((4, 3))
object_points[:, :2] = square_size * np.array([[-0.5, -0.5], [-0.5, 0.5], [0.5, 0.5], [0.5, -0.5]])
# 使用PnP算法解算相对位姿
success, rvec, tvec = cv2.solvePnP(object_points, image_points, K, None)

if success:
    # 计算重投影误差
    reproj_error = 0
    for i in range(len(object_points)):
        img_point, _ = cv2.projectPoints(object_points[i], rvec, tvec, K, None)
        error = cv2.norm(image_points[i], img_point.squeeze(), cv2.NORM_L2) ** 2
        reproj_error += error

    # 打印重投影误差
    print("角点检测误差本身由于不合理投影带来的重投影误差: ", reproj_error/len(object_points), "像素")
else:
    print("PnP failed to converge")

R, _ = cv2.Rodrigues(rvec)
pose = np.column_stack((R, tvec))
arc = rotation_matrix_to_euler_angles(R)

theta = np.linalg.norm(rvec) / np.pi * 180
print("每个角点", pix_rev, "像素的误差引起的最大角度误差为：")
print(theta, "°")
print("写成欧拉角形式：")
print("r=", arc[0][0], "°,p=", arc[1][0], "°,y=", arc[2][0], "°")
